// .mtl file structure:
//   newmtl: Declare new material group, with name
//   Ka: ambient color
//   Kd: diffuse color
//   Ks: specular color
//   Ke: emissive color
//   Ns: brightness
//   d, Tr: Alpha
//   sharpness: sharpness
//   illum： illumination (0~10)
//   Ni : optical density
//   Tf: transmission filter (r，g，b)
//   refl: reflection properties
//   map_Ka: ambient map
//   map_Kd: diffuse map
//   map_Ks: specular map

#include "ObjLoader.h"

#include <kgl.h>

#include <memory>
#include <fstream>

#include <stb_image.h>

ObjLoader::ObjLoader(const std::string &fileName) {
    load_obj(fileName);
    kp::Ref<kp::MatteMaterial> default_mat;
    SHARED_ASSET_TRY_GET("default_mat", kp::MakeRef<kp::MatteMaterial>(), default_mat)
//    auto uv_tex_mat = SHARED_ASSET_ADD_GET("uv_tex_mat", kp::MakeRef<kp::MatteMaterial>(kp::UVTexture::Create()));
    // after loading .mtl, check materials in mat_list exist in mat_map or not
    for (auto&[index, mtl_name] : mtl_list) {
        auto tex_path = mtl_map[mtl_name].map_Kd;
        if (!tex_path.empty()) {
            kp::Ref<kp::ImageTexture> img_tex;
            SHARED_ASSET_GET(tex_path, img_tex);
            if (nullptr == img_tex) {
                int w, h, ch;
                auto stbi_del = [](unsigned char *data) { stbi_image_free(data); };
                kp::Scope<unsigned char[], decltype(stbi_del)> dataptr(
                        stbi_load(tex_path.c_str(), &w, &h, &ch, 0), stbi_del);
                kp::Ref<kp::Mipmap<kp::Spectrum>> mpmp;
                // get mipmap
                SHARED_ASSET_TRY_GET(tex_path, kp::Mipmap<kp::Spectrum>::Create(w, h, ch, dataptr), mpmp)
                // get img_tex
                SHARED_ASSET_TRY_GET(tex_path, kp::MakeRef<kp::ImageTexture>(mpmp), img_tex)
            }
            kp::Ref<kp::MatteMaterial> cur_mat;
            SHARED_ASSET_TRY_GET(tex_path, kp::MakeRef<kp::MatteMaterial>(img_tex), cur_mat)
            mesh.mat_map.insert({mtl_name, cur_mat});
        } else {
            kp::Ref<kp::MatteMaterial> cur_mat;
            SHARED_ASSET_TRY_GET(mtl_name, kp::MakeRef<kp::MatteMaterial>(
                    kp::MakeRef<kp::ConstantTexture>(mtl_map[mtl_name].Kd)), cur_mat)
            mesh.mat_map.insert({mtl_name, cur_mat});
        }
    }

}

std::vector<kp::Ref<kp::Primitive>> ObjLoader::get(kp::Ref<kp::Transform> o2w) {
    kgl::Log::Info("model info>> v:{0}, t:{1}, n:{2}, f:{3}",
                   mesh.vertices.size(),
                   mesh.uvs.size(),
                   mesh.normals.size(),
                   mesh.faceindices.size() / 3);

    int nTriangles = (int) mesh.faceindices.size() / 3;
    int nVertices = (int) mesh.vertices.size();
    auto tris = kp::TriangleMesh::Create(o2w, o2w->inverseRef(), false,
                                         nTriangles, nVertices,
                                         mesh.faceindices.data(),
                                         mesh.vertices.data(),
                                         mesh.normals.data(),
                                         nullptr,
                                         mesh.uvs.data());
    std::vector<kp::Ref<kp::Primitive>> ret;
    kp::Ref<kp::MatteMaterial> default_mat;
    SHARED_ASSET_TRY_GET("default_mat", kp::MakeRef<kp::MatteMaterial>(), default_mat)
    int list_index = 0;
    for (int i = 0; i < tris.size(); i++) {
        if (list_index < mtl_list.size() && i >= mtl_list[list_index].first) {
            ++list_index;
        }
        auto &mat_ref = mesh.mat_map[mtl_list[list_index - 1].second];
        if (nullptr == mat_ref)mat_ref = default_mat;
        ret.push_back(kp::MakeRef<kp::GeometryPrimitive>(mat_ref, tris[i]));
    }
    return ret;
}

void ObjLoader::load_obj(const std::string &fileName) {
// Get mtllib
    std::ifstream in(fileName);
    if (!in.good()) {
        kgl::Log::Error("loading obj:'{0}' file is not good", fileName);
        exit(0);
    }
    char buffer[256];
    char mtl_name[256];
    char mtllib[256];
    kp::Point3f ver;
    kp::Normal nor;
    kp::Point2f uv;
    kp::Face f1, f2, f3;

    while (!in.getline(buffer, 255).eof()) {
        if (buffer[0] == 'f' && buffer[1] == ' ') {
            int nt = sscanf(buffer, "f %d/%d/%d %d/%d/%d %d/%d/%d",
                            &f1.v, &f1.t, &f1.n,
                            &f2.v, &f2.t, &f2.n,
                            &f3.v, &f3.t, &f3.n
            );
            mesh.faceindices.push_back(f1 - kp::Point3i(1, 1, 1));
            mesh.faceindices.push_back(f2 - kp::Point3i(1, 1, 1));
            mesh.faceindices.push_back(f3 - kp::Point3i(1, 1, 1));
            if (nt != 9) {
                kgl::Log::Error("ERROR: I don't know the format of that FaceMtl");
            }
        } else if (buffer[0] == 'v' && buffer[1] == ' ') {
            if (sscanf(buffer, "v %f %f %f", &ver.x, &ver.y, &ver.z) == 3) {
                mesh.vertices.push_back(ver);
            } else {
                kgl::Log::Error("ERROR: vertex not in wanted format in OBJLoader");
            }
        } else if (buffer[0] == 'v' && buffer[1] == 't') {
            if (sscanf(buffer, "vt %f %f", &uv.x, &uv.y) == 2) {
                mesh.uvs.push_back(uv);
            } else {
                kgl::Log::Error("ERROR: texCoord not in wanted format in OBJLoader");
            }
        } else if (buffer[0] == 'v' && buffer[1] == 'n') {
            if (sscanf(buffer, "vn %f %f %f", &nor.x, &nor.y, &nor.z) == 3) {
                mesh.normals.push_back(nor);
            } else {
                kgl::Log::Error("ERROR: normal not in wanted format in OBJLoader");
            }
        } else if (buffer[0] == 'u' && buffer[1] == 's') {
            if (sscanf(buffer, "usemtl %s", mtl_name) == 1) {
                // append start_index with mat_name into mat_list, whatever existed or not
//                kgl::Log::Info("start index {0} will use mtl: {1}", mesh.faceindices.size() / 3, mtl_name);
                mtl_list.emplace_back(mesh.faceindices.size() / 3, mtl_name);
            }
        } else if (buffer[0] == 'm' && buffer[1] == 't') {
            if (sscanf(buffer, "mtllib %s", mtllib) == 1) {
                // load .mtl file in other thread
                auto mtl_path = kgl::fs::replace_file(fileName, mtllib);
//                kgl::Log::Info("mtllib path: {0}", mtl_path);
                load_mtl(mtl_path);
            } else {
                kgl::Log::Error("ERROR: mtllib path not in wanted format");
            }
        }
    }
}

void ObjLoader::load_mtl(const std::string &fileName) {
    std::ifstream in(fileName);
    if (!in.good()) {
        kgl::Log::Error("loading obj:({0}) file is not good", fileName);
        exit(0);
    }
    char buffer[256];
    char mtl_name[256]; // current mtl name
    char tex_file[256];
    while (!in.getline(buffer, 255).eof()) {
        if (buffer[0] == 'n' && buffer[1] == 'e') {
            // newmtl mtl_name (TODO name without space now)
            sscanf(buffer, "newmtl %s", mtl_name);
            mtl_map.insert({mtl_name, {}});
        } else if (buffer[0] == 'm' && buffer[1] == 'a' && buffer[2] == 'p') {
            if (buffer[4] == 'K' && buffer[5] == 'd') {
                sscanf(buffer, "map_Kd %s", tex_file);
                auto tex_path = kgl::fs::replace_file(fileName, tex_file);
                if (mtl_map.find(mtl_name) != mtl_map.end()) {
                    mtl_map[mtl_name].map_Kd = tex_path;
                } else {
                    kgl::Log::Error("map_Kd: mtl name '{0}' not found.", mtl_name);
                }
            } else if (buffer[4] == 'K' && buffer[5] == 's') {
                sscanf(buffer, "map_Ks %s", tex_file);
                auto tex_path = kgl::fs::replace_file(fileName, tex_file);
                if (mtl_map.find(mtl_name) != mtl_map.end()) {
                    mtl_map[mtl_name].map_Ks = tex_path;
                } else {
                    kgl::Log::Error("map_Ks: mtl name '{0}' not found.", mtl_name);
                }
            }
        } else if (buffer[0] == 'N') {
            if (buffer[1] == 's') {
                float tmp_ns;
                sscanf(buffer, "Ns %f", &tmp_ns);
                if (mtl_map.find(mtl_name) != mtl_map.end()) {
                    mtl_map[mtl_name].Ns = tmp_ns;
                } else {
                    kgl::Log::Error("Ns: mtl name '{0}' not found.", mtl_name);
                }
            } else if (buffer[1] == 'i') {
                float tmp_ni;
                sscanf(buffer, "Ni %f", &tmp_ni);
                if (mtl_map.find(mtl_name) != mtl_map.end()) {
                    mtl_map[mtl_name].Ni = tmp_ni;
                } else {
                    kgl::Log::Error("Ni: mtl name '{0}' not found.", mtl_name);
                }
            }
        } else if (buffer[0] == 'K') {
            if (buffer[1] == 'a') {
                kp::Spectrum tmp_ka;
                int inputs = sscanf(buffer, "Ka %f %f %f", &tmp_ka.r, &tmp_ka.g, &tmp_ka.b);
                if (inputs == 3) {
                    if (mtl_map.find(mtl_name) != mtl_map.end()) {
                        mtl_map[mtl_name].Ka = tmp_ka;
                    } else {
                        kgl::Log::Error("Ka: mtl name '{0}' not found.", mtl_name);
                    }
                } else {
                    kgl::Log::Error("Ka: not float3.", mtl_name);
                }
            } else if (buffer[1] == 'd') {
                kp::Spectrum tmp_kd;
                int inputs = sscanf(buffer, "Kd %f %f %f", &tmp_kd.r, &tmp_kd.g, &tmp_kd.b);
                if (inputs == 3) {
                    if (mtl_map.find(mtl_name) != mtl_map.end()) {
                        mtl_map[mtl_name].Kd = tmp_kd;
                    } else {
                        kgl::Log::Error("Kd: mtl name '{0}' not found.", mtl_name);
                    }
                } else {
                    kgl::Log::Error("Kd: not float3.", mtl_name);
                }
            } else if (buffer[1] == 's') {
                kp::Spectrum tmp_ks;
                int inputs = sscanf(buffer, "Ks %f %f %f", &tmp_ks.r, &tmp_ks.g, &tmp_ks.b);
                if (inputs == 3) {
                    if (mtl_map.find(mtl_name) != mtl_map.end()) {
                        mtl_map[mtl_name].Ks = tmp_ks;
                    } else {
                        kgl::Log::Error("Ks: mtl name '{0}' not found.", mtl_name);
                    }
                } else {
                    kgl::Log::Error("Ks: not float3.", mtl_name);
                }
            } else if (buffer[1] == 'e') {
                kp::Spectrum tmp_ke;
                int inputs = sscanf(buffer, "Ke %f %f %f", &tmp_ke.r, &tmp_ke.g, &tmp_ke.b);
                if (inputs == 3) {
                    if (mtl_map.find(mtl_name) != mtl_map.end()) {
                        mtl_map[mtl_name].Ke = tmp_ke;
                    } else {
                        kgl::Log::Error("Ke: mtl name '{0}' not found.", mtl_name);
                    }
                } else {
                    kgl::Log::Error("Ke: not float3.", mtl_name);
                }
            }
        } else if ((buffer[0] == 'd' && buffer[1] == ' ') ||
                   (buffer[0] == 'T' && buffer[1] == 'r')) {
            float tmp_d;
            if (buffer[0] == 'd')
                sscanf(buffer, "d %f", &tmp_d);
            else
                sscanf(buffer, "Tr %f", &tmp_d);
            if (mtl_map.find(mtl_name) != mtl_map.end()) {
                mtl_map[mtl_name].d_ = tmp_d;
            } else {
                kgl::Log::Error("d: mtl name '{0}' not found.", mtl_name);
            }
        } else if (buffer[0] == 'i' && buffer[1] == 'l') {
            int tmp_illum;
            sscanf(buffer, "illum %d", &tmp_illum);
            if (mtl_map.find(mtl_name) != mtl_map.end()) {
                mtl_map[mtl_name].illum = tmp_illum;
            } else {
                kgl::Log::Error("illum: mtl name '{0}' not found.", mtl_name);
            }
        }
    }
}
